#! python
#-*- coding: utf-8 -*- 


from abc import ABC, abstractmethod
import numpy as np
from typing import TypeAlias, Union
from numpy.typing import NDArray

Real: TypeAlias = np.float64
Integer: TypeAlias = Union[int, np.int64]

class Distribution(ABC):
    """概率分布抽象基类"""
    def __init__(self, dimension: int = 1):
        self._dimension = dimension
        
    @abstractmethod
    def random(self, size: int) -> NDArray[Real]:
        """ 生成该分布的样本 """
        pass

    @abstractmethod
    def logp(self, samples: NDArray[Real]) -> NDArray[Real]:
        """ 计算样本的概率密度函数的对数 
        
        参数:
            samples: 输入样本数组，形状为(dimension, n_samples)
            
        返回:
            对数概率密度值
        """
        pass


class UniformDistribution(Distribution):
    """均匀分布实现"""
    def __init__(self, low: float = 0.0, high: float = 1.0, dimension: int = 1):
        super().__init__(dimension)
        self.low = low
        self.high = high
        
    def random(self, size: int) -> NDArray[Real]:
        """生成均匀分布随机数"""
        return np.random.uniform(self.low, self.high, (self._dimension, size))
        
    def logp(self, samples: NDArray[Real]) -> NDArray[Real]:
        """计算均匀分布的对数概率密度"""
        in_range = (samples >= self.low) & (samples <= self.high)
        log_prob = np.full(samples.shape[1], -np.inf)
        log_prob[np.all(in_range, axis=0)] = -np.log(self.high - self.low) * self._dimension
        return log_prob

class RealGaussian(Distribution):
    def __init__(self, mean: Union[float, NDArray[Real]] = 0.0, std: float = 1.0):
        """初始化实正态分布
        
        参数:
            dimension: 分布维度(可选，默认根据mean自动推断)
            mean: 均值 (默认0.0)，可以是标量或数组
            std: 标准差 (默认1.0)
        """
        # 处理mean参数
        if np.isscalar(mean) or isinstance(mean, float) or isinstance(mean, int):
            dimension: int = 1
        else:
            dimension: int = int(np.prod(mean.shape))
        super().__init__(dimension)
        self._mean = np.asarray(mean)
        self._std = std
        
    def random(self, size:int) -> NDArray[Real]:
        """ 随机生成高斯随机数 
        >>> p = RealGaussian()
        >>> data = p.random(10000)
        >>> data.shape
        (1, 10000)
        >>> data.dtype
        dtype('float64')
        >>> np.abs(np.mean(data)) < 0.2 or np.mean(data) # 验证均值
        np.True_
        >>> np.abs(np.std(data) - 1) < 0.2  # 验证标准差
        np.True_
        >>> p2 = RealGaussian(1.0, 2.0)
        >>> data = p2.random(10000)
        >>> data.shape
        (1, 10000)
        >>> data.dtype
        dtype('float64')
        >>> np.all(np.abs(np.mean(data, axis=1) - 1.0) < 0.5) or np.mean(data, axis=1)  # 验证均值
        np.True_
        >>> np.all(np.abs(np.std(data, axis=1) - 2.0) < 0.5)  # 验证标准差
        np.True_
        """
        from numpy import random
        dimension: int = self._dimension
        return self._mean.reshape(-1, 1) + self._std * random.randn(dimension, size)
        
    def logp(self, samples: NDArray[Real]) -> NDArray[Real]:
        """ 计算样本的概率密度函数的对数 
        
        参数:
            samples: 输入样本数组，形状为(dimension, n_samples)
            
        返回:
            每个样本的对数概率密度值数组，形状为(n_samples,)
            
        示例:
        >>> p = RealGaussian()
        >>> samples = p.random(5)
        >>> log_probs = p.logp(samples)
        >>> log_probs.shape  # 返回每个样本的对数概率
        (5,)
        >>> not np.all(log_probs < 0)  # 对数概率应为负数
        False
        >>> isinstance(log_probs[0], Real)  # 元素类型检查
        True
        >>> p2 = RealGaussian(1.0, 2.0)
        >>> samples = p2.random(5)
        >>> log_probs = p2.logp(samples)
        >>> log_probs.shape  # 返回每个样本的对数概率
        (5,)
        >>> not np.all(log_probs < 0)  # 对数概率应为负数
        False
        >>> isinstance(log_probs[0], Real)  # 元素类型检查
        True
        """
        if samples.ndim != 2 or samples.shape[0] != self._dimension:
            raise ValueError(f"输入样本形状应为({self._dimension}, n), 但得到{samples.shape}")
            
        # 正态分布的对数概率密度
        # log p(x) = -0.5 * ((x-μ)^2/σ^2 + d*log(2πσ^2))
        z = (samples - self._mean.reshape(-1, 1)) / self._std
        log_std = np.log(self._std)
        return -0.5 * (np.sum(z**2, axis=0) + self._dimension * (np.log(2 * np.pi) + 2 * log_std))

    def cdf(self, x: NDArray[Real]) -> NDArray[Real]:
        """计算累积分布函数
        
        参数:
            x: 输入点，形状为(dimension, n_points)
            
        返回:
            每个维度上的累积概率，形状为(dimension, n_points)
            
        示例:
        >>> p = RealGaussian(0, 1)
        >>> x = np.array([[0.0]])
        >>> np.abs(p.cdf(x) - 0.5) < 1e-10  # CDF at mean should be 0.5
        array([[ True]])
        >>> p = RealGaussian([0, 1], 1)
        >>> x = np.array([[0.0], [1.0]])
        >>> probs = p.cdf(x)
        >>> np.all(np.abs(probs - 0.5) < 1e-10)  # CDF at mean should be 0.5
        True
        """
        if x.ndim != 2 or x.shape[0] != self._dimension:
            raise ValueError(f"输入形状应为({self._dimension}, n), 但得到{x.shape}")
            
        # 标准化
        z = (x - self._mean.reshape(-1, 1)) / self._std
        # 使用误差函数计算累积概率
        from scipy import special
        return 0.5 * (1 + special.erf(z / np.sqrt(2)))
    
    def ppf(self, q: NDArray[Real]) -> NDArray[Real]:
        """计算累积分布函数的逆（分位数函数）
        
        参数:
            q: 概率值，取值范围[0,1]，形状为(dimension, n_points)
            
        返回:
            对应的分位数，形状为(dimension, n_points)
            
        示例:
        >>> p = RealGaussian(0, 1)
        >>> q = np.array([[0.5]])
        >>> np.abs(p.ppf(q)) < 1e-10  # median should be at mean
        array([[ True]])
        >>> p = RealGaussian([0, 1], 1)
        >>> q = np.array([[0.5], [0.5]])
        >>> x = p.ppf(q)
        >>> np.allclose(x, np.array([[0], [1]]))  # should return means
        True
        """
        if q.ndim != 2 or q.shape[0] != self._dimension:
            raise ValueError(f"输入形状应为({self._dimension}, n), 但得到{q.shape}")
            
        if np.any((q < 0) | (q > 1)):
            raise ValueError("概率值必须在[0,1]范围内")
            
        # 使用逆误差函数计算分位数
        from scipy.special import erfinv
        return self._mean.reshape(-1, 1) + self._std * np.sqrt(2) * erfinv(2 * q - 1)

    def interval(self, alpha: float) -> NDArray[Real]:
        """计算置信区间
        
        参数:
            alpha: 显著性水平，如0.05表示95%置信区间
            
        返回:
            每个维度的区间边界，形状为(dimension, 2)
            
        示例:
        >>> p = RealGaussian(0, 1)
        >>> interval = p.interval(0.05)
        >>> np.allclose(interval, np.array([[-1.96, 1.96]]))  # 95% CI
        True
        >>> p = RealGaussian([0, 1], 1)
        >>> interval = p.interval(0.05)
        >>> np.allclose(interval[0], np.array([-1.96, 1.96]))  # 95% CI for dim 1
        True
        >>> np.allclose(interval[1], np.array([-0.96, 2.96]))  # 95% CI for dim 2
        True
        """
        if alpha <= 0 or alpha >= 1:
            raise ValueError("alpha必须在(0,1)范围内")
            
        q = np.array([[alpha/2, 1-alpha/2]] * self._dimension)
        return self.ppf(q)
